import cv2
import pydicom
from pydicom.pixel_data_handlers.util import apply_voi_lut
import numpy as np


def _lin_stretch_img(img, low_prc, high_prc, do_ignore_minmax=True):
    """
    Apply linear "stretch" - low_prc percentile goes to 0,
    and high_prc percentile goes to 255.
    The result is clipped to [0, 255] and converted to np.uint8

    Additional feature:
    When computing high and low percentiles, ignore the minimum and maximum intensities (assumed to be outliers).
    """
    # For ignoring the outliers, replace them with the median value
    if do_ignore_minmax:
        tmp_img = img.copy()
        med = np.median(img)  # Compute median
        tmp_img[img == img.min()] = med
        tmp_img[img == img.max()] = med
    else:
        tmp_img = img

    lo, hi = np.percentile(tmp_img, (low_prc, high_prc))  # Example: 1% - Low percentile, 99% - High percentile

    if lo == hi:
        return np.full(img.shape, 128, np.uint8)  # Protection: return gray image if lo = hi.

    stretch_img = (img.astype(float) - lo) * (255 / (hi - lo))  # Linear stretch: lo goes to 0, hi to 255.
    stretch_img = stretch_img.clip(0, 255).astype(np.uint8)  # Clip range to [0, 255] and convert to uint8
    return stretch_img


def convert(dcm_in_file: str, png_out_file: str):
    """
    转换
    :param dcm_in_file: 输入dcm
    :param png_out_file: 输出png
    :return:
    """
    # https://www.visus.com/fileadmin/content/pictures/Downloads/JiveX_DICOME_Viewer/case1.zip
    ds = pydicom.read_file(dcm_in_file)  # read dicom image
    img = ds.pixel_array  # get image array

    # https://pydicom.github.io/pydicom/stable/old/working_with_pixel_data.html#voi-lut-or-windowing-operation
    # Apply "VOI gray scale transformations":
    img = apply_voi_lut(img, ds, index=0)

    # Apply "linear stretching" (lower percentile 0.1 goes to 0, and percentile 99.9 to 255).
    img = _lin_stretch_img(img, 0.1, 99.9)

    # https://dicom.innolitics.com/ciods/rt-dose/image-pixel/00280004
    if ds[0x0028, 0x0004].value == 'MONOCHROME1':
        # Invert polarity if Photometric Interpretation is 'MONOCHROME1'
        img = 255 - img

    cv2.imwrite(png_out_file, img)
